import json

# 需求：
# 可以一次性上传一整个文件夹的文件，并启动工作流生成微调预料，并最后获取下来，存成一个json。
# 宏观思路
# 1. 先获取文件夹下的文件。
# 2. 通过上传文件的api上传文件，得到文件的id列表
# 3. 把文件id列表生成满足执行工作流的api这个传参形式，调用执行工作流的 api
# 4. 获取工作流结果，对结果进行解析，解析成json格式，保存成json文件。


api_key = 'app-WqiQCDEUM0Hlm8Jvpj0jGtNu'

import requests
from uuid import uuid4 # 生成随机字符串
import os
# 封装一个上传文件的函数
def upload_file(local_path):
    url = 'http://116.62.130.29/v1/files/upload'
    headers = {'Authorization': f'Bearer {api_key}' }
    files = {'file': (f'{str(uuid4())}.pdf', open(local_path, 'rb'), 'application/pdf')}
    data = {'user': 'abc-123'}
    response = requests.post(url, headers=headers, files=files, json=data)
    return response.json()['id']

# 封装一个上传一个文件夹下所有文件的函数
def uploads(dir_path):
    # dir_path 文件夹路径
    ids = []
    for file_name in os.listdir(dir_path):
        local_path = os.path.join(dir_path, file_name)
        id = upload_file(local_path)
        ids.append(id)
    return ids #返回所有上传文件的id


# 封装一个执行工作流的函数
def run_workflow(ids):
    url = 'http://116.62.130.29/v1/workflows/run'
    headers = {
        'Authorization': f'Bearer {api_key}',
        'Content-Type': 'application/json'
    }
    files = []
    for id in ids:
        file = {
            "transfer_method": "local_file",
            "upload_file_id": id,
            "type": "document"
        }
        files.append(file)

    data = {
        "inputs": {"files":files},
        "user": "abc-123"
    }
    res = requests.post(url,headers=headers,json=data)
    return res.json()


# 封装一个解析结果的函数
def parse_json_and_save(d):
    d = d['data']['outputs']['text']
    d = '[' + d + ']'
    d = json.loads(d)
    with open('output1.json', 'w+',encoding='utf-8') as f:
        json.dump(d, f, indent=4, ensure_ascii=False)

# 封装一个最外层的函数
def do():
    ids = uploads("docs")
    print(ids)
    d = run_workflow(ids)
    parse_json_and_save(d)

if __name__ == '__main__':
    do()
